#!/usr/bin/env python
# coding: utf-8

"""
predict_rddhw.py
create by cuiyu 2025-02-23
说明：
  * 识别图片文件，将识别结果保存在图像相同的目录下
    + .json：为 labelme 格式的标注文件
    + .aaby: 为 aaby 格式（实际也是 json 文件）的识别结果文件
  * 依赖文件：
    + classes_rddhw_labelme.json : labelme 标注格式的类型定义文件 labelx 需要有前导的数字
 
用法:
  predict-rddhy.py

更新:
2025-02-23:
"""

# 移除未使用的导入项
# import datetime
import json
import sys
from pathlib import Path
from loguru import logger as log
from ultralytics import YOLO
import ultralytics.engine.results as ulresults
import rmfile_predict
from bhclasses_define import BHClasses
from runtimer import runtimer_create
import torch

_APP_FILE = Path(sys.argv[0])
_APP_ROOT = _APP_FILE.parent

# import os
# os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:1024"


def _classes_file_load(file: Path) -> bool:
    """
    加载并解析 classes-rddhw-labelme.json : labelme 标注格式的类型定义文件
    """
    if not file.exists():
        return False
    
    json_str = file.read_text(encoding="utf-8")
    global _G_MODEL_CLASSES
    try:
        _G_MODEL_CLASSES = json.loads(json_str)
        # G_LABELS = dict(zip(labels_.values(), labels_.keys()))
    except:
        return False
    
    return True

def _result_save(result: ulresults.Results, img_path: Path, decimals: int=5) -> bool:
    """
    将单个预测结果转换为 aaby JSON格式，并保存为 aaby 文件

    参数:
        :results: (ultralytics.engine.results.Results): 单个预测结果对象
        :img_path: Path 本次识别对应的图像的全路径名
        :decimals: int: 保留小数位数，0=整数，5=保留5位小数
    返回:
        :bool: False=没有识别出病害，True=识别出病害
    """
    # 1. 保存 labelme json 文件
    hasBH = _write_labelme_json(result, img_path)


    if hasBH:
        # 2. 保存 .aaby json 文件
        kresults_dict = result.summary(decimals=decimals)
        mask_dcell = _get_mask_cell(result, 50)
        aaby_dict: dict = {
            "mask_dcell": mask_dcell.tolist(),
            "kresults": kresults_dict,
        }
        _write_json_file(aaby_dict, img_path, ".aaby")

        # 3. 保存 .cell json 文件 (将 aaby_dict 中的 segments 删除)
        for kr in aaby_dict["kresults"]:
            kr.pop("segments", "")
        _write_json_file(aaby_dict, img_path, ".cell")

    return hasBH

def _write_json_file(json_dict: dict, img_path: Path, extname: str) -> bool:
    """
    将单个预测结果转换为 aaby JSON格式，并保存为 aaby 文件

    参数:
        :results: (ultralytics.engine.results.Results): 单个预测结果对象
        :img_path: Path 本次识别对应的图像的全路径名

    返回:
        :bool: False=没有识别出病害，True=识别出病害
    """
    img_path = Path(img_path)
    # json_str = json.dumps(json_dict, indent=4, ensure_ascii=False)
    json_str = json.dumps(json_dict, ensure_ascii=False)
    json_file = img_path.with_suffix(extname)
    json_file.write_text(json_str, encoding="utf-8")
    # with open(json_file, "w", encoding="utf-8") as f:
    #     f.write(json_str)
    return True



def _write_labelme_json(result: ulresults.Results, img_path: Path) -> bool:
    """
    将单个预测结果转换为 Labelme JSON格式，并保存为JSON文件

    参数:
        :results: (ultralytics.engine.results.Results): 单个预测结果对象
        :img_path: Path 本次识别对应的图像的全路径名

    返回:
        :bool: False=没有识别出病害，True=识别出病害
    """

    img_path = Path(img_path)
    labelme_json = {
        "version": "4.5.6",
        "flags": {},
        "shapes": [],
        "imagePath": str(img_path.name),
        "imageData": None,
        "imageHeight": result.orig_shape[0],
        "imageWidth": result.orig_shape[1]
    }

    # 检查是否有检测结果
    if (result.boxes is None) or (result.masks is None):
        return False
    if (len(result.boxes) == 0) or (len(result.masks)) == 0:
        return False
    xys = result.masks.xy
    if (xys is None) or (len(xys) == 0):
        return False

    # 遍历每个检测框
    for box, points, cls in zip(result.boxes.data, xys, result.boxes.cls):
        shape = {
            "label":  BHClasses.getName(int(cls)),  # 使用类别映射表获取类别名称 _G_MODEL_CLASSES[str(int(cls))],  # 使用类别映射表获取类别名称
            "points": [list(map(float, point)) for point in points],  # 多边形点列表
            "group_id": None,
            "shape_type": "polygon",
            "flags": {}
        }
        if len(shape["points"])<=3:
            continue
        labelme_json["shapes"].append(shape)
    
    if (len(labelme_json["shapes"]) == 0):
        return False
    
    json_str = json.dumps(labelme_json, indent=4, ensure_ascii=False)
    json_file = img_path.with_suffix(".json")
    json_file.write_text(json_str, encoding="utf-8")
    return True

# 保存 cell 格式的函数定义
def _get_mask_cell(result: ulresults.Results, step: int) -> torch.Tensor:
    """ 
    从 Result 中获取聚会后的 dcell masks 
    参数：
        :result:     YOLO 的预测结果
        :step:       池化步长
        :decimals:   浮点数保存的精度，小数点后位数

    """
    
    #==== 已通过YOLOv11获取分割结果 results[0] 是第一个图像的结果
    masks = result.masks.data  # 所有对象的掩膜, masks.data 才是 tensor, masks.data.shape = (9, 1024, 2048)
    classes = result.boxes.cls+1  # 各对象的类别ID + 1 shape = (9)
    obj_num = len(classes)
    # img_size = result.orig_shape  # 原始图像尺寸 (2000, 4096)
    # dcell_size = (20, 40)  # dcell 一般为 (20, 40) 个 = 200✕400厘米 

    #==== 将 masks 池化为 (20, 40) 的尺寸
    # maxpool = torch.nn.AdaptiveMaxPool2d(dcell_size)
    maxpool = torch.nn.MaxPool2d((55, 59), stride=51, padding=0)  # 确保输出为 (20, 40)
    masks_pool: torch.Tensor = maxpool(masks.data.unsqueeze(0)) # masks_pool.sharp =  (1, 9, 20, 40)
    # for i, (mask, cls_id) in enumerate(zip(masks_pool[0], classes)):
    #     count = torch.sum(mask)  # 统计 True 的个数
    #     print(f"cls_id = {cls_id-1}, mask.shape = {masks.shape}, count = {count}")


    #==== 将各对象的掩膜与类别ID相乘 * 1000 （后三位留给 对象ID）
    classes_ = classes.unsqueeze(0).unsqueeze(-1).unsqueeze(-1)  # 形状变为 (1, 9, 1, 1)
    masks_cls = masks_pool * classes_ * 1000 # shape = (1, 9, 20, 40)
    # for i, (mask, cls_id) in enumerate(zip(masks_cls[0], classes)):
    #     count = torch.sum(mask)  # 统计 True 的个数
    #     print(f"cls_id = {cls_id-1}, masks_cls.shape = {masks_cls.shape}, count = {count}")


    #==== 将各对象的掩膜与对象ID相乘 （后三位留给 类别ID）
    obj_ids = torch.arange(1, obj_num+1, dtype=masks.dtype, device=masks.device)  # 生成1-9的ID序列
    # obj_ids = obj_ids.unsqueeze(0).unsqueeze(-1).unsqueeze(-1)  # 形状变为 (1, 9, 1, 1) 这个也可以，会自动扩展维度
    obj_ids = obj_ids.view(1, obj_num, 1, 1).expand_as(masks_pool)  # 维度扩展 (1, 9, 20, 40) 与 masks_pool 维度一致 
    masks_obj = masks_pool * obj_ids    # shape = (1, 9, 20, 40)

    # 聚合所有对象的掩膜
    masks_all = masks_cls + masks_obj  # shape = (1, 9, 20, 40)
    masks_dcell, _ =  masks_all.max(dim=1, keepdim=True)  # shape = (1, 1, 20, 40), dtype=double
    masks_dcell = masks_dcell.squeeze().to(dtype=torch.uint8)  # # 默认移除所有大小为1的维度 shape = (20, 40), dtype=uint8
    return masks_dcell



def _predict_file(model:YOLO, file:Path, imgsz, conf:float=0.1, iou:float=0.1) -> int:
    """
    对单个文件进行预测

    参数：
        :model:  加载的YOLO模型
        :file:   待识别的图像文件
        :imgze:  图像尺寸
        :conf:   置信度
        :iou:    IOU 阈值，控制重叠输出，越小输出的重叠的矩形框越少
    返回：
        :int:    0=没有识别出病害，1=识别出病害
    """
    r_cnt = False
    results = model.predict(source=file, verbose=False, imgsz=imgsz, retina_masks=False, conf=conf, iou=iou)
    for r in results:
        r_cnt = _result_save(r, file) or r_cnt
    return 1 if r_cnt else 0
    

def _predictDir(model, in_dir: Path, imgsz, conf=0.1, iou=0.1) -> tuple[int, int]:
    """ 对指定目录下的图像文件进行预测

    参数：
        :model:  加载的YOLO模型
        :in_dir: 需要处理的图像目录路径
        :imgze:  图像尺寸
        :conf:   置信度
        :iou: float;   IOU 阈值，控制重叠输出，越小输出的重叠的矩形框越少
    返回：
        :(int, int): (处理的文件总数, 识别出的文件总数)
    """
    # 1. 按照文件名排序返回 in_dir 下的 .json 文件列表，准备进度窗
    img_files = [Path(file) for file in sorted(in_dir.rglob("*.jpg"), key=lambda file: str(file))]

    # 2. 循环处理每个文件，并显示进度
    bh_sum = 0
    file_sum = len(img_files)
    n_len = len(str(file_sum))
    print(f"开始推理：目录 = {in_dir}, 文件数 = {file_sum}")
    for i, file in enumerate(img_files):
        n = i+1
        # 打印进度
        print(f'{n:{n_len}} / {file_sum} {float(n)/file_sum*100:.2f}%  {str(file.name)}', end='\r')
        bh_sum += _predict_file(model, file, imgsz, conf, iou)
    print("")
    print(f"处理文件数：{file_sum}, 发现病害文件数：{bh_sum}")
    return (file_sum, bh_sum)
    

def predictRddhw(model_file: Path, data_root: Path, imgsz, conf=0.1, iou=0.1)->tuple[int, int]:
    """
    对指定目录下的图像文件进行预测，一个子目录一个子目录的分批进行
    
    Args:
      :model_file: str;  模型文件全路径名
      :data_root: str;  数据目录，包含若干个子目录，每个子目录下都是待识别的图像文件
      :imgsz: int;  图像尺寸
      :conf: float;  置信度
      :iou: float;   IOU 阈值，控制重叠输出，越小输出的重叠的矩形框越少
    
    Returns:
      :int, int: 处理的文件总数, 发现病害的文件总数
    """
    # log start
    if not model_file.exists():
        log.error(f"模型文件不存在: {str(model_file.absolute())}")
        return (0, 0)
    if not data_root.exists():
        log.error(f"数据目录不存在: {str(data_root)}")
        return (0, 0)
    

    # 1. Load a model
    model = YOLO(model_file, task="segment")

    # 删除 data_root 下的所有 .json 和 .aaby 文件
    with runtimer_create("删除文件"):
        rmfile_predict.run(data_root, "*.[ja][sa][ob][ny]")  # 
        # rmfile_predict.run(data_root, "*.json")  # [ja][sa][ob][ny]
    
    # 2. 按照目录名排序返回 data_root 的下一级子目录列表
    img_dirs = sorted(
        [child for child in data_root.iterdir() if child.is_dir()], 
        key=lambda x: x.name.lower()
    )
    
    # 3. 循环处理每个目录，并显示进度
    file_sum = hasBH_sum = 0
    for dir in img_dirs:
        files, bhs = _predictDir(model, dir, imgsz, conf, iou)
        file_sum += files
        hasBH_sum += bhs

    # 4. 在 data_root 下创建一个名为 aaby 的文件
    aaby_file = data_root / "aaby"
    aaby_file.touch(exist_ok=True)

    # 返回处理的文件总数和识别出病害的文件总数
    return (file_sum, hasBH_sum)




def _logInit(baseDir: Path) -> None:
    log.add(baseDir/"predict-labelme.log", rotation="500 MB", encoding='utf-8')  

def _main():
    from runtimer import runtimer_create
    
    model_file = Path("/aaby/netdisk/mydata/ai-data/model/rddhw/rddhw-20250404/weights/best.engine")   # 有结果输出的模型 
    # data_root = Path("/aaby/netdisk/mydata/ai-data/dataset/priv/rddhw/自动标注数据/img01")
    data_root = Path("/aaby/netdisk/mydata/ai-data/dataset/priv/rddhw/评估数据/内蒙/C316二级路+右幅+-上行-2车道-20240727142816")
    imgsz = 2048
    conf = 0.01
    iou = 0.01

    with runtimer_create("图片识别") as timer:
        file_sum, bh_sum = predictRddhw(model_file, data_root, imgsz, conf, iou)
    print(f"识别完成，文件总数：{file_sum}, 识别速度：{float(file_sum)/timer.elapsed * 60:.2f} 张/分, 识别病害文件总数：{bh_sum}")

    # log start

if __name__ == "__main__":
    _logInit(_APP_ROOT / "log/")
    _main()
